# app.py
from flask import flash, request, redirect, url_for, Flask, render_template, send_from_directory, session
from flask_sqlalchemy import SQLAlchemy
from flask_login import LoginManager, UserMixin, login_user, login_required, logout_user, current_user
from flask_wtf import FlaskForm
from wtforms import StringField, PasswordField, SubmitField
from wtforms.validators import DataRequired, Length, EqualTo
from werkzeug.security import generate_password_hash, check_password_hash
import os
from datetime import datetime
from werkzeug.utils import secure_filename
import zipfile
from parse_data import ParseData
import shutil

app = Flask(__name__)
app.config['SECRET_KEY'] = 'dypkxt'  # 必须设置一个随机密钥
app.config['SQLALCHEMY_DATABASE_URI'] = 'mysql+pymysql://admin:Wing1Q2W#E@120.55.126.68/crawler'  # 替换为你的MySQL信息
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False

app.config['UPLOAD_FOLDER'] = 'static/uploads'
app.config['UNZIP_FOLDER'] = 'static/unzip/'
app.config['ALLOWED_EXTENSIONS'] = {'zip'}
app.config['EXCEL_PATH'] = 'static/excel/'
db = SQLAlchemy(app)
login_manager = LoginManager(app)
login_manager.login_view = 'login'

# 确保上传目录存在
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
os.makedirs(app.config['UNZIP_FOLDER'], exist_ok=True)

def allowed_file(filename):
    return '.' in filename and \
        filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']



# 用户模型
class User(UserMixin, db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(50), unique=True, nullable=False)
    password_hash = db.Column(db.String(128), nullable=False)

    def set_password(self, password):
        self.password_hash = generate_password_hash(password, method='pbkdf2:sha256')

    def check_password(self, password):
        return check_password_hash(self.password_hash, password)

# 注册表单
class RegistrationForm(FlaskForm):
    username = StringField('用户名', validators=[DataRequired(), Length(min=4, max=50)])
    password = PasswordField('密码', validators=[DataRequired(), Length(min=6)])
    confirm_password = PasswordField('确认密码', validators=[DataRequired(), EqualTo('password')])
    submit = SubmitField('注册')

# 登录表单
class LoginForm(FlaskForm):
    username = StringField('用户名', validators=[DataRequired()])
    password = PasswordField('密码', validators=[DataRequired()])
    submit = SubmitField('登录')

@login_manager.user_loader
def load_user(user_id):
    return User.query.get(int(user_id))

@app.route('/')
def home():
    return render_template('home.html')

@app.route('/register', methods=['GET', 'POST'])
def register():
    form = RegistrationForm()
    if form.validate_on_submit():
        user = User(username=form.username.data)
        user.set_password(form.password.data)
        db.session.add(user)
        db.session.commit()
        flash('注册成功！请登录', 'success')
        return redirect(url_for('login'))
    return render_template('register.html', form=form)

@app.route('/login', methods=['GET', 'POST'])
def login():
    form = LoginForm()
    if form.validate_on_submit():
        user = User.query.filter_by(username=form.username.data).first()
        if user and user.check_password(form.password.data):
            login_user(user)
            return redirect(url_for('dashboard'))
        else:
            flash('用户名或密码错误', 'danger')
    return render_template('login.html', form=form)

@app.route('/dashboard')
@login_required
def dashboard():
    return render_template('upload.html', username=current_user.username)


@app.route('/logout')
@login_required
def logout():
    logout_user()
    return redirect(url_for('home'))


@app.route('/upload', methods=['POST'])
@login_required
def upload():
    # 获取表单数据
    title = request.form.get('title')
    company = request.form.get('company')
    date = request.form.get('date')
    remarks = request.form.get('remarks')

    # 获取上传文件
    if 'file' not in request.files:
        return redirect(request.url)
    file = request.files['file']

    if file.filename == '':
        return redirect(request.url)

    if file and allowed_file(file.filename):
        # 生成安全的文件名
        filename = secure_filename(file.filename)
        # 添加时间戳防止重名
        new_filename = f"{datetime.now().strftime('%Y%m%d%H%M%S')}_{filename}"
        save_path = os.path.join(app.config['UPLOAD_FOLDER'], new_filename)
        file.save(save_path)

        # 这里可以添加数据库存储逻辑（示例仅打印）
        unzip_path = app.config['UNZIP_FOLDER']
        with zipfile.ZipFile(save_path, 'r') as zip_ref:
            zip_ref.extractall(unzip_path)
        dir = unzip_path + filename.replace(".zip", "")
        json_path = app.config['EXCEL_PATH'] + f"{datetime.now().strftime('%Y%m%d%H%M%S')}_all_item.json"
        excel_path = app.config['EXCEL_PATH'] + f"{datetime.now().strftime('%Y%m%d%H%M%S')}_all_item.xlsx"
        result_path = app.config['EXCEL_PATH'] + f"{datetime.now().strftime('%Y%m%d%H%M%S')}_all_result.xlsx"
        try:
            p = ParseData(xls_path_dir=dir, json_path=json_path, excel_path=excel_path, result_path=result_path)
            p.main()
        except Exception as e:
            print(e)
        finally:
            os.remove(save_path)
            # shutil.rmtree(unzip_path)

        filenames = [excel_path.replace(app.config['EXCEL_PATH'], ""), result_path.replace(app.config['EXCEL_PATH'], "")]
        session['filenames'] = filenames
        return filenames,200

    return '无效的文件类型'

@login_required
@app.route('/success')
def success():
    filenames = session.pop('filenames', None)
    return render_template('download.html', filenames=filenames)

@login_required
@app.route('/download/<filename>')
def download_file(filename):
    """提供文件下载"""
    return send_from_directory(
        app.config['EXCEL_PATH'],
        filename,
        as_attachment=True
    )


if __name__ == '__main__':
    # with app.app_context():
    #     db.create_all()  # 创建数据库表（首次运行时）
    app.run(host="0.0.0.0", port=5000, debug=True)